import argparse
import numpy as np


def load_train_config(para, layers_num=None):
    # 设置参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--DRR_train_path', type=str,
                        default=para + "data/CT/均匀采样/train/标准正位/DRR",
                        help='train data location')
    parser.add_argument('--train_json_path', type=str,
                        default=para + "data/CT/均匀采样/train/标准正位/label.json",
                        help='train data_label location')
    parser.add_argument('--DRR_test_path', type=str,
                        default=para + "data/CT/均匀采样/test/标准正位/DRR",
                        help='tester data location')
    parser.add_argument('--test_json_path', type=str,
                        default=para + "data/CT/均匀采样/test/标准正位/label.json",
                        help='tester data_label location')
    parser.add_argument('--end_epoch', type=int, default=50000)
    parser.add_argument('--is_resume', type=bool, default=False, help='whether loading the checkpoint')
    parser.add_argument('--CT_path', type=str, default=para + 'CT_data', help='loading the CT voxel')
    parser.add_argument('--mode', type=str, default='reg', help='training mode')
    parser.add_argument('--pre_model_address', type=str,
                        default=para + "data/history/cla10000/model/6D_model_loss1.9206700325012207.pth",
                        help='loading the pretrain model')
    parser.add_argument('--mlp_num', type=int, default=128, help='the mlp number for model,only useful for UBasic')
    parser.add_argument('--conv_scale', type=int, default=1,
                        help='the conv structure for model,only useful for UBasic')
    parser.add_argument('--layers_num', type=int, default=layers_num,
                        help='the conv depth for model,only useful for UBasic')
    parser.add_argument('--voxel_size', type=np.ndarray, default=np.array([180.0, 180.0, 107.0]),
                        help='the size of CT voxel')
    parser.add_argument('--interval_num_train', type=np.ndarray, default=np.array([20, 20, 4]),
                        help='the sample interval of CT voxel when training')
    parser.add_argument('--interval_num_test', type=np.ndarray, default=np.array([20, 20, 4]),
                        help='the sample interval of CT voxel when testing')
    parser.add_argument('--rot_cen', type=np.ndarray, default=np.array([89.82421875, 89.82421875, 53]),
                        help='the rotation center of CT voxel')
    parser.add_argument('--d_s2c', type=int, default=400,
                        help='the distance between light source and the center of CT voxel')
    args = parser.parse_args()
    para_dict = {
        'DRR_train_path': args.DRR_train_path,
        'train_json_path': args.train_json_path,
        'DRR_test_path': args.DRR_test_path,
        'test_json_path': args.test_json_path,
        'end_epoch': args.end_epoch,
        'is_resume': args.is_resume,
        'CT_path': args.CT_path,
        'mode': args.mode,
        'pre_model_address': args.pre_model_address,
        'mlp_num': args.mlp_num,
        'conv_scale': args.conv_scale,
        'layers_num': args.layers_num,
        'voxel_size': args.voxel_size,
        'interval_num_train': args.interval_num_train,
        'interval_num_test': args.interval_num_test,
        'rot_cen': args.rot_cen,
        'd_s2c': args.d_s2c
    }
    return para_dict
